package com.mjf.spark.day08

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{DataFrame, Dataset, Encoder, Encoders, SparkSession, TypedColumn}

/**
 * 自定义UDAF(强类型  主要应用在DSL风格的DS查询)
 */
object SparkSQL06_UDAF {
  def main(args: Array[String]): Unit = {

    // 创建SparkConf配置文件对象
    val conf: SparkConf = new SparkConf().setAppName("SparkSQL").setMaster("local[*]")

    // 创建SparkSQL执行的入口对象  SparkSession
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()

    import spark.implicits._

    // 读取json文件，创建DF
    val df: DataFrame = spark.read.json("D:\\code\\spark\\input\\test.json")

    // 创建自定义函数对象
    val myAvgNew: myAvgNew = new myAvgNew

    // 注意：如果是自定义UDAF的强类型，没有办法应用SQL风格DF的查询
//    spark.udf.register("myAvgNew", myAvgNew)

    // 将df转换为ds
    val ds: Dataset[User06] = df.as[User06]

    // 将自定义函数对象转换为查询列
    val col: TypedColumn[User06, Double] = myAvgNew.toColumn

    // 在进行查询时，会将查询出来的记录（User06类型）交给自定义的函数处理
    ds.select(col).show()

    // 释放资源
    spark.close()

  }
}

// 注意：如果从内存中获取数据，spark可以知道数据类型具体是什么，如果是数字，默认作为Int处理
// 但是从文件中读取的数字，不能确定是什么类型，所以用bigint接收，可以和Long类型转换，但是和Int不能进行转换
// 输入类型的样例类
case class User06(name:String, age:Long)

// 缓存类型的样例类
case class AgeBuffer(var sum:Long, var count:Long)

// 自定义UDAF函数（强类型）
// @tparam IN  输入数据类型
// @tparam BUF 缓存数据类型
// @tparam OUT 输出结果数据类型
class myAvgNew extends Aggregator[User06, AgeBuffer, Double]{

  // 缓存初始化
  override def zero: AgeBuffer = {
    AgeBuffer(0L, 0L)
  }

  // 对当前分区内数据进行聚合
  override def reduce(b: AgeBuffer, a: User06): AgeBuffer = {
    b.sum += a.age
    b.count += 1
    b
  }

  // 分区间合并
  override def merge(b1: AgeBuffer, b2: AgeBuffer): AgeBuffer = {
    b1.count += b2.count
    b1.sum += b2.sum
    b1
  }

  // 返回计算结果
  override def finish(reduction: AgeBuffer): Double = {
    reduction.sum.toDouble / reduction.count
  }

  // DataSet的编码以及解码器，用于进行序列化，固定写法
  // 用户自定义类型  product     系统值类型，根据具体类型进行选择即可
  override def bufferEncoder: Encoder[AgeBuffer] = {
    Encoders.product
  }

  //
  override def outputEncoder: Encoder[Double] = {
    Encoders.scalaDouble
  }
}